from collections import Callable
from typing import List

import torch


def apply(*objects: List, transform: Callable) -> List:
    return [transform(o) for o in objects]


def cuda(*objects: List) -> List:
    return apply(*objects, transform=lambda x: x.cuda())

def todevice(device, *objects: List) -> List:
    return apply(*objects, transform=lambda x: x.to(device))

def clone_unsqueeze_cuda(*objects: List) -> List:
    return apply(*objects, transform=lambda x: x.clone().unsqueeze(0).cuda())

def squeeze_permute_cpu_np(*objects: List) -> List:
    return apply(
        *objects,
        transform=lambda x: x.squeeze(0).permute((1, 2, 0)).cpu().numpy(),
    )

def permute_cpu_np(*objects: List) -> List:
    return apply(
        *objects,
        transform=lambda x: x.permute((1, 2, 0)).cpu().numpy(),
    )

if __name__ == '__main__':
    x = torch.rand((3,256,256)).cuda()
    y = torch.rand((3,256,256)).cuda()
    x,y = clone_unsqueeze_cuda(x,y)
    print(x.shape)
    x,y= squeeze_permute_cpu_np(x,y)
    print(x.shape)
